{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9TV7IYeqifSv"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=ByZjmtFgB_Y5)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kZRlD4utdPuX"
      },
      "outputs": [],
      "source": [
        "%install '.package(url: \"https://github.com/tensorflow/swift-models\", .branch(\"tensorflow-0.11\"))' Datasets ImageClassificationModels\n",
        "print(\"\\u{001B}[2J\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tRIJp_4m_Afz"
      },
      "outputs": [],
      "source": [
        "// #@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "// you may not use this file except in compliance with the License.\n",
        "// You may obtain a copy of the License at\n",
        "//\n",
        "// https://www.apache.org/licenses/LICENSE-2.0\n",
        "//\n",
        "// Unless required by applicable law or agreed to in writing, software\n",
        "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "// See the License for the specific language governing permissions and\n",
        "// limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sI1ZtrdiA4aY"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        " <td><a target=\"_blank\" href=\"https://www.tensorflow.org/swift/tutorials/introducing_x10\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org で表示</a></td>\n",
        " <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/swift/tutorials/introducing_x10.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab で実行</a></td>\n",
        " <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/swift/tutorials/introducing_x10.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub でソースを表示</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8sa42_NblqRE"
      },
      "source": [
        "# X10 入門\n",
        "\n",
        "Swift for TensorFlow は、デフォルトでは Eager execution を使用してテンソル演算を実行します。これは迅速なイテレーションを可能にしていますが、機械学習モデルのトレーニングで最もパフォーマンスが高いオプションというわけではありません。\n",
        "\n",
        "[X10 テンソルライブラリ](https://github.com/tensorflow/swift-apis/blob/master/Documentation/X10/API_GUIDE.md)は、テンソルトレースと [XLA コンパイラ](https://www.tensorflow.org/xla)を利用して、Swift for TensorFlow に高性能なバックエンドを追加します。このチュートリアルでは、X10 を紹介し、トレーニングループを更新して GPU や TPU で実行するプロセスを案内します。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7MpNcIwIIy8"
      },
      "source": [
        "## Eager テンソルと X10 テンソル"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lM9dRji7IIy8"
      },
      "source": [
        "Swift for TensorFlow の加速化計算は、テンソルの型を通して実行されます。テンソルは様々な演算に参加することができるため、機械学習モデルの基本的なビルディングブロックです。\n",
        "\n",
        "デフォルトでは、Tensor は Eager execution を使用して演算ごとに計算を実行します。各 Tensor には、接続されているハードウェアとそれに使用されているバックエンドを説明する Device が関連づけられています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHRTNQJo1TxT"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "import Foundation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FCMWR11NIIy-"
      },
      "outputs": [],
      "source": [
        "let eagerTensor1 = Tensor([0.0, 1.0, 2.0])\n",
        "let eagerTensor2 = Tensor([1.5, 2.5, 3.5])\n",
        "let eagerTensorSum = eagerTensor1 + eagerTensor2\n",
        "eagerTensorSum"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1qad9_yMYf6F"
      },
      "outputs": [],
      "source": [
        "eagerTensor1.device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrlMNOinIIy_"
      },
      "source": [
        "このノートブックを GPU 対応のインスタンスで実行している場合は、上記のデバイス説明内にハードウェアが反映されているはずです。Eager ランタイムは TPU をサポートしていないため、TPU をアクセラレータとして使用している場合には、ハードウェアターゲットとして CPU が使用されていることが分かります。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eoyLeSQVIIy9"
      },
      "source": [
        "Tensor を作成する際に代替デバイスを指定して、Eager モードのデフォルトの デバイスを上書きすることができます。これが X10 バックエンドを使用した計算の実行をオプトインする方法です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NrRQhQaHaJm9"
      },
      "outputs": [],
      "source": [
        "let x10Tensor1 = Tensor([0.0, 1.0, 2.0], on: Device.defaultXLA)\n",
        "let x10Tensor2 = Tensor([1.5, 2.5, 3.5], on: Device.defaultXLA)\n",
        "let x10TensorSum = x10Tensor1 + x10Tensor2\n",
        "x10TensorSum"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VbqeudQCaqwv"
      },
      "outputs": [],
      "source": [
        "x10Tensor1.device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mIbIOW0HIIzA"
      },
      "source": [
        "これを GPU 対応のインスタンスで実行している場合は、X10 テンソルのデバイス内にアクセラレータのリストが表示されます。これを Eager execution ではなく TPU 対応のインスタンスで実行している場合は、計算にそのデバイスを使用していることが確認できるはずです。X10は Swift for TensorFlow で TPU を活用する方法なのです。\n",
        "\n",
        "デフォルトの Eager デバイスと X10 デバイスは、システム上の最初のアクセラレータを使用しようとします。GPU が接続されている場合、最初に利用可能な GPU を使用します。TPU が存在する場合は、X10 はデフォルトで最初の TPU コアを使用します。アクセラレータが見つからない、またはサポートされていない場合、デフォルトのデバイスは CPU にフォールバックします。\n",
        "\n",
        "デフォルトの Eager や XLA デバイスを越えて、Device 内に固有のハードウェアターゲットとバックエンドターゲットを提供することができます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "De59VwJ35SvG"
      },
      "outputs": [],
      "source": [
        "// let tpu1 = Device(kind: .TPU, ordinal: 1, backend: .XLA)\n",
        "// let tpuTensor1 = Tensor([0.0, 1.0, 2.0], on: tpu1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rU0WY_sJodio"
      },
      "source": [
        "## Eager モードモデルをトレーニングする\n",
        "\n",
        "デフォルトの Eager execution モードを使用し、モデルを設定してトレーニングする方法を見てみましょう。この例では [swift-models リポジトリ](https://github.com/tensorflow/swift-models)および MNIST 手書き数字分類データセットから単純な LeNet-5 モデルを使用します。\n",
        "\n",
        "まず、MNIST データセットをセットアップしてダウンロードします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqXILiXhq-iM"
      },
      "outputs": [],
      "source": [
        "import Datasets\n",
        "\n",
        "let epochCount = 5\n",
        "let batchSize = 128\n",
        "let dataset = MNIST(batchSize: batchSize)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pMewsl0VgnJ"
      },
      "source": [
        "次に、モデルとオプティマイザの設定を行います。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fEAEyUExXT3I"
      },
      "outputs": [],
      "source": [
        "import ImageClassificationModels\n",
        "\n",
        "var eagerModel = LeNet()\n",
        "var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cwNDwzS2QgS1"
      },
      "source": [
        "ここで、基本的な進捗状況の追跡とレポートを実装します。すべての中間統計はトレーニングが実行されるデバイスと同じデバイス上でテンソルとして保持され、`scalarized()` はレポート作成中にのみ呼び出されます。これで不必要に遅延テンソルがマテリアライズされるのを回避できるため、後で X10 を使用する場合に特に重要になります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eNzOdly3QY_P"
      },
      "outputs": [],
      "source": [
        "struct Statistics {\n",
        "    var correctGuessCount = Tensor<Int32>(0, on: Device.default)\n",
        "    var totalGuessCount = Tensor<Int32>(0, on: Device.default)\n",
        "    var totalLoss = Tensor<Float>(0, on: Device.default)\n",
        "    var batches: Int = 0\n",
        "    var accuracy: Float { \n",
        "        Float(correctGuessCount.scalarized()) / Float(totalGuessCount.scalarized()) * 100 \n",
        "    } \n",
        "    var averageLoss: Float { totalLoss.scalarized() / Float(batches) }\n",
        "\n",
        "    init(on device: Device = Device.default) {\n",
        "        correctGuessCount = Tensor<Int32>(0, on: device)\n",
        "        totalGuessCount = Tensor<Int32>(0, on: device)\n",
        "        totalLoss = Tensor<Float>(0, on: device)\n",
        "    }\n",
        "\n",
        "    mutating func update(logits: Tensor<Float>, labels: Tensor<Int32>, loss: Tensor<Float>) {\n",
        "        let correct = logits.argmax(squeezingAxis: 1) .== labels\n",
        "        correctGuessCount += Tensor<Int32>(correct).sum()\n",
        "        totalGuessCount += Int32(labels.shape[0])\n",
        "        totalLoss += loss\n",
        "        batches += 1\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w3lmTRCWT5sS"
      },
      "source": [
        "最後に、トレーニングループを使用してモデルを 5 エポック分、実行します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W9bUsiOxVf_v"
      },
      "outputs": [],
      "source": [
        "print(\"Beginning training...\")\n",
        "\n",
        "for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {\n",
        "    let start = Date()\n",
        "    var trainStats = Statistics()\n",
        "    var testStats = Statistics()\n",
        "    \n",
        "    Context.local.learningPhase = .training\n",
        "    for batch in batches {\n",
        "        let (images, labels) = (batch.data, batch.label)\n",
        "        let 𝛁model = TensorFlow.gradient(at: eagerModel) { eagerModel -> Tensor<Float> in\n",
        "            let ŷ = eagerModel(images)\n",
        "            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "            trainStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "            return loss\n",
        "        }\n",
        "        eagerOptimizer.update(&eagerModel, along: 𝛁model)\n",
        "    }\n",
        "\n",
        "    Context.local.learningPhase = .inference\n",
        "    for batch in dataset.validation {\n",
        "        let (images, labels) = (batch.data, batch.label)\n",
        "        let ŷ = eagerModel(images)\n",
        "        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "        testStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "    }\n",
        "\n",
        "    print(\n",
        "        \"\"\"\n",
        "        [Epoch \\(epoch)] \\\n",
        "        Training Loss: \\(String(format: \"%.3f\", trainStats.averageLoss)), \\\n",
        "        Training Accuracy: \\(trainStats.correctGuessCount)/\\(trainStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", trainStats.accuracy))%), \\\n",
        "        Test Loss: \\(String(format: \"%.3f\", testStats.averageLoss)), \\\n",
        "        Test Accuracy: \\(testStats.correctGuessCount)/\\(testStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", testStats.accuracy))%) \\\n",
        "        seconds per epoch: \\(String(format: \"%.1f\", Date().timeIntervalSince(start)))\n",
        "        \"\"\")\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ED0HGZW2gWY"
      },
      "source": [
        "ご覧のように、モデルは期待通りにトレーニングされ、検証セットに対する精度はエポックごとに向上しました。Swift for TensorFlow モデルはこのようにして定義され、Eager execution を使用して実行されています。ここで、X10 を活用するためにはどのような修正が必要かを考えてみましょう。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Te7sNNx9c_am"
      },
      "source": [
        "## X10 モデルをトレーニングする\n",
        "\n",
        "データセット、モデル、オプティマイザには、Eager execution デバイスのデフォルトで初期化されるテンソルが含まれています。X10 で作業をするには、これらのテンソルを X10 デバイスに移動する必要があります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaN7fM-lAe7r"
      },
      "outputs": [],
      "source": [
        "let device = Device.defaultXLA\n",
        "device"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cJfSg0wDAgC7"
      },
      "source": [
        "データセットについては、バッチがトレーニングループで処理されている時点で移動させ、Eager execution モデルからデータセットを再利用できるようにします。\n",
        "\n",
        "モデルとオプティマイザについては、Eager execution デバイス上の内部テンソルで初期化して、X10 デバイスに移動させます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpcOByipc75O"
      },
      "outputs": [],
      "source": [
        "var x10Model = LeNet()\n",
        "x10Model.move(to: device)\n",
        "\n",
        "var x10Optimizer = SGD(for: x10Model, learningRate: 0.1)\n",
        "x10Optimizer = SGD(copying: x10Optimizer, to: device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQvza3dUXlr0"
      },
      "source": [
        "トレーニングループに必要な修正は、いくつかの特定の時点で行います。まず最初に、トレーニングデータのバッチを X10 デバイスに移動させる必要があります。これは各バッチを取得する際に `Tensor(copy:to:)` を介して行います。\n",
        "\n",
        "次の修正は、トレーニングループの実行中にトレースをどこで切断するかを示すことです。X10 はコードに必要なテンソル計算をトレースし、そのトレースの最適化された表現をジャストインタイムでコンパイルすることによって動作します。トレーニングループの場合は同じ演算を何度も繰り返すことになるので、トレース、コンパイル、再利用には理想的なセクションです。\n",
        "\n",
        "Tensor からの値を明示的に要求するコード（これらは通常 `.scalars` または `.scalarized()` なので目立ちます）がない場合、X10 はすべてのループイテレーションをまとめてコンパイルしようとします。これを防ぎ、特定の時点でトレースを切断するために、オプティマイザがモデル重みを更新して検証中に損失と精度を取得した後で、明示的な `LazyTensorBarrier()` を配置します。これにより、トレーニングループの各ステップと検証中の推論の各バッチから成る 2 つの再利用トレースが作成されます。\n",
        "\n",
        "これらの修正によって、以下のトレーニングループが発生します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XrZee8n3Y17_"
      },
      "outputs": [],
      "source": [
        "print(\"Beginning training...\")\n",
        "\n",
        "for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {\n",
        "    let start = Date()\n",
        "    var trainStats = Statistics(on: device)\n",
        "    var testStats = Statistics(on: device)\n",
        "    \n",
        "    Context.local.learningPhase = .training\n",
        "    for batch in batches {\n",
        "        let (eagerImages, eagerLabels) = (batch.data, batch.label)\n",
        "        let images = Tensor(copying: eagerImages, to: device)\n",
        "        let labels = Tensor(copying: eagerLabels, to: device)\n",
        "        let 𝛁model = TensorFlow.gradient(at: x10Model) { x10Model -> Tensor<Float> in\n",
        "            let ŷ = x10Model(images)\n",
        "            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "            trainStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "            return loss\n",
        "        }\n",
        "        x10Optimizer.update(&x10Model, along: 𝛁model)\n",
        "        LazyTensorBarrier()\n",
        "    }\n",
        "\n",
        "    Context.local.learningPhase = .inference\n",
        "    for batch in dataset.validation {\n",
        "        let (eagerImages, eagerLabels) = (batch.data, batch.label)\n",
        "        let images = Tensor(copying: eagerImages, to: device)\n",
        "        let labels = Tensor(copying: eagerLabels, to: device)\n",
        "        let ŷ = x10Model(images)\n",
        "        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "        LazyTensorBarrier()\n",
        "        testStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "    }\n",
        "\n",
        "    print(\n",
        "        \"\"\"\n",
        "        [Epoch \\(epoch)] \\\n",
        "        Training Loss: \\(String(format: \"%.3f\", trainStats.averageLoss)), \\\n",
        "        Training Accuracy: \\(trainStats.correctGuessCount)/\\(trainStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", trainStats.accuracy))%), \\\n",
        "        Test Loss: \\(String(format: \"%.3f\", testStats.averageLoss)), \\\n",
        "        Test Accuracy: \\(testStats.correctGuessCount)/\\(testStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", testStats.accuracy))%) \\\n",
        "        seconds per epoch: \\(String(format: \"%.1f\", Date().timeIntervalSince(start)))\n",
        "        \"\"\")\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qej_Z6V3mZnG"
      },
      "source": [
        "X10 バックエンドを使用したモデルトレーニングは、以前の Eager execution モデルの場合と全く同じように進行するはずです。最初のバッチの前と最初のエポックの終わりに遅延が生じる可能性がありますが、これはその時点で固有のトレースをジャストインタイムでコンパイルしているためです。アクセラレータを接続して実行している場合は、その時点以降のトレーニングが Eager モードよりも速く進んでいることが分かるはずです。\n",
        "\n",
        "初期トレースのコンパイル時間と高速なスループットのトレードオフがありますが、大部分の機械学習モデルでは、繰り返し演算によるスループットの増加がコンパイルのオーバーヘッドを相殺するよりも大きくなるはずです。実際に複数のトレーニングケースにおいて、X10 を使用することによってスループットが 4 倍以上向上しています。\n",
        "\n",
        "先に述べたとおり、X10 を使用すると TPU での作業ができるようになるだけでなく、作業が容易で、Swift for TensorFlow モデルのアクセラレータのすべてのクラスを使用できるようになります。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "introducing_x10.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Swift",
      "name": "swift"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
